import os
import argparse
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

from core.freq4 import EPUBWordCounter, parse_target_files


def main():
    parser = argparse.ArgumentParser(description='统计EPUB文件中英文单词频率并识别生词（多线程版本）')
    parser.add_argument('directory', help='包含EPUB文件的目录路径')
    parser.add_argument('-n', '--new-books-dir', help='新EPUB文件目录路径（用于识别生词）')
    parser.add_argument('-f', '--target-files', help='指定只解析新书中的特定文件列表，用英文逗号分隔（如: Chapter_Five.xhtml,Section0023.xhtml）')
    parser.add_argument('-o', '--output', default='word_frequency.xlsx',
                        help='输出Excel文件路径 (默认: word_frequency.xlsx)')
    parser.add_argument('--new-words-dir', default='output',
                        help='新词输出目录路径 (默认: output)')
    parser.add_argument('--known-words-output', default='known_words.json',
                        help='已知单词列表输出文件路径 (默认: known_words.json)')
    parser.add_argument('-t', '--top-n', type=int, default=1000,
                        help='输出前N个最常出现的单词 (默认: 1000)')
    parser.add_argument('-c', '--context-words', type=int, default=12,
                        help='例句上下文单词数量 (默认: 12)')
    parser.add_argument('-r', '--frequency-threshold', type=float, default=0.01,
                        help='已知单词频率百分比阈值 (默认: 0.01%%)')
    parser.add_argument('--no-cross-file-deduplication', action='store_true',
                        help='禁用跨文件去重（默认启用）')
    parser.add_argument('--no-cache', action='store_true',
                        help='禁用缓存（默认启用）')
    parser.add_argument('--threads', type=int, default=4,
                        help='处理新书时使用的线程数 (默认: 4)')

    args = parser.parse_args()

    # 解析目标文件列表
    target_files = parse_target_files(args.target_files)

    # 创建输出目录
    os.makedirs(args.new_words_dir, exist_ok=True)

    # 创建计数器实例
    counter = EPUBWordCounter()

    # 处理目录中的EPUB文件
    if not counter.process_directory(args.directory, args.frequency_threshold, not args.no_cache):
        return

    # 保存已知单词列表
    counter.save_known_words(args.known_words_output, args.frequency_threshold)

    # 保存主要结果到Excel
    if not counter.save_to_excel(args.output, args.top_n, args.frequency_threshold):
        return

    # 如果有新书目录，处理新书中的生词
    if args.new_books_dir:
        if not os.path.exists(args.new_books_dir):
            logger.error(f"新书目录不存在: {args.new_books_dir}")
            return

        # 获取新书目录中的所有EPUB文件
        new_epub_files = [f for f in os.listdir(args.new_books_dir)
                          if f.lower().endswith('.epub')]

        if not new_epub_files:
            logger.error("在新书目录中未找到EPUB文件")
            return

        logger.info(f"找到 {len(new_epub_files)} 本新书，使用 {args.threads} 个线程处理")

        # 使用线程池处理多本新书
        with ThreadPoolExecutor(max_workers=args.threads) as executor:
            futures = []

            for epub_file in new_epub_files:
                epub_path = os.path.join(args.new_books_dir, epub_file)
                future = executor.submit(
                    counter.process_new_book,
                    epub_path,
                    args.new_words_dir,
                    args.context_words,
                    not args.no_cross_file_deduplication,
                    target_files,
                    args.frequency_threshold
                )
                futures.append(future)

            # 等待所有任务完成
            success_count = 0
            fail_count = 0
            for future in as_completed(futures):
                success, book_path = future.result()
                if success:
                    success_count += 1
                else:
                    fail_count += 1

            logger.info(f"新书处理完成: 成功 {success_count} 本, 失败 {fail_count} 本")

    # 打印统计信息
    logger.info(f"处理完成! 共处理 {counter.processed_files} 个文件")
    logger.info(f"总单词数: {counter.total_words}")
    logger.info(f"唯一单词数: {len(counter.word_counter)}")
    logger.info(f"已知单词数（频率≥{args.frequency_threshold}%）: {len(counter.known_words)}")
    logger.info(f"主要结果已保存到: {args.output}")

    if args.new_books_dir:
        logger.info(f"生词结果已保存到目录: {args.new_words_dir}")


if __name__ == "__main__":
    main()